#include "VariantStats.h"

#include <algorithm>
#include <iostream>
#include <numeric>
#include <functional>
#include <limits.h>
#include "Variant.h"

using namespace std;
using namespace vcf;

/**
 * This function reads variants of interest from a variants file name.
 * It fills variantMap with one TargetVariant per variant of interest
 */
void read_variants(string interesting_variant_filename, 
		   map<GeneLocation, TargetVariant*>& variantMap) {
  int numvariants = INT_MAX;
  ifstream interesting_variant_file;
  interesting_variant_file.open(interesting_variant_filename.c_str(), ifstream::in);

  string oneVariant;

  while (std::getline(interesting_variant_file, oneVariant)) {
    vector<string> fields = split(oneVariant, "\t");
    chromosome c = fields[0];
    int basepair = atoi(fields[1].c_str());
    TargetVariant *t = new TargetVariant(c, basepair);
    variantMap.insert(pair<GeneLocation, TargetVariant*>(t->location, t));
  }
  interesting_variant_file.close();
}

/**
 * This function reads in aggregation buckets from a file.  The format is 
 *
 * <chromosome><tab><base pair range>\n
 * 
 * Lines beginning with '#' are ignored.  It's possible to only
 * specify a chromosome, too.
 */
void read_aggregation_buckets(string aggregation_bucket_filename, vector<GeneLocation>& aggregation_buckets) {
  ifstream bucket_file;
  bucket_file.open(aggregation_bucket_filename.c_str(), ifstream::in);

  string one_bucket;
  while (std::getline(bucket_file, one_bucket)) {
    if (one_bucket[0] == '#') {
      continue;
    }
    chromosome c;
    int bp_start = NOT_A_RANGE, bp_end = NOT_A_RANGE;
    if (one_bucket.find('\t') != -1) {
      vector<string> fields = split(one_bucket, "\t");
      c = fields[0];
      if (fields[1].find('-') != -1) {
	vector<string> bp_range = split(fields[1],'-');
	bp_start = atoi(bp_range[0].c_str());
	bp_end = atoi(bp_range[1].c_str());
      } else {
	bp_start = atoi(fields[1].c_str());
	bp_end = NOT_A_RANGE;
      }
    } else {
      c = one_bucket.c_str();
    }
    GeneLocation g(c, bp_start, bp_end);
    aggregation_buckets.push_back(g);
  }
}

/** 
 * This function actually parses vcf files and stores aggregate data
 * from all VCF files for the same variant together */
void accumulate_read_counts_for_vcf_file(const char* vcf_filename, 
					 map<GeneLocation, TargetVariant*>& variantMap,
					 string chromosome_filter) {
  VariantCallFile variantFile;

  variantFile.open(*(new string(vcf_filename)));
  if (!variantFile.is_open()) {
    cerr << "Could not open file " << vcf_filename << endl;
    return;
  }

  Variant var(variantFile);
  string chromosome;
  int bp;
  int variantNumber = 0;
  while (variantFile.getNextVariantNoParse(var, chromosome, bp)) {
    variantNumber++;
    if (chromosome_filter != "") {
      if (chromosome != chromosome_filter) {
	continue;
      }
    }

    var.parse(variantFile.line);
    if (var.filter == "LowQual") {
      continue;
    }
    
    map<string, vector<string> >::iterator infoFlagIter;
    infoFlagIter = var.info.find("refseq.changesAA");
    if (infoFlagIter == var.info.end()) {
      continue;
    }
    bool changesAA = ((*infoFlagIter).second[0] == "true");

    if (!changesAA) {
      continue;
    }

    GeneLocation g(chromosome, bp);

    map<GeneLocation, TargetVariant*>::iterator iter = variantMap.find(g);

    if (iter == variantMap.end()) {
      pair<GeneLocation, TargetVariant*> p(g, new TargetVariant(chromosome, bp));
      variantMap.insert(p);
      iter = variantMap.find(g);
    } else {
      // Seeing duplicate lines in our file, so skipping.
      continue;
    }

    pair<GeneLocation, TargetVariant*> p = *iter;
    VariantAndCounts* vc = new VariantAndCounts();
    // Copy the alleles for this variant into our own TargetVariant structure.
    vc->alleles = var.alleles;
    TargetVariant* t = p.second;
    t->variant_stats.push_back(vc);

    int ref_count = 0, alt_count = 0;

    map<string, map<string, vector<string> > >::iterator it;
    // Iterate over all samples, collecting information on the number
    // of hetero/homozygous samples, and the read counts of each
    // allele.
    for (it = var.samples.begin(); it != var.samples.end(); ++it) {
      vector<string> gtConfidence = it->second["GQ"];
      if (gtConfidence.size() != 1) {
	cout << "skipping " << variantNumber << " because confidence size was: " << gtConfidence.size() << endl;
	continue;
      }

      if (atoi(gtConfidence[0].c_str()) < 99) {
	continue;
      }
      vector<string> ar = it->second["AD"];
      vector<string> depth = it->second["DP"];
      int readdepth = atoi(depth[0].c_str());
      int alt = atoi(ar[0].c_str());
      int ref = atoi(ar[1].c_str());
      
      if (alt + ref > 2 * readdepth) {
	continue;
      }

      vector<string> genotypeInfo = it->second["GT"];
      assert(genotypeInfo.size() == 1);
      // Add genotype information to our VariantAndCounts structure
      string genotype_string = genotypeInfo[0];
      map<genotype, int>& genotypeCounts = vc->genotypeCounts;
      map<genotype, int>::iterator findResult = genotypeCounts.find(genotype_string);
      if (findResult == genotypeCounts.end()) {
	pair<genotype, int> p;
	p.first = genotype_string;
	p.second = 1;
	genotypeCounts.insert(p);
      } else {
	(*findResult).second++;
      }

      // Now match allele read counts to zygosity and store.
      vector<string> alleleReadCounts = it->second["AD"];
      assert(alleleReadCounts.size() == 2);

      zygosity zyg;
      if (genotype_string[0] == genotype_string[2]) {
	zyg = HOM;
      } else {
	zyg = HET;
      }

      int allele1Reads = atoi(alleleReadCounts[0].c_str());
      int allele2Reads = atoi(alleleReadCounts[1].c_str());

      if (zyg == HET || genotype_string[0] != '0') {
	vc->ref_reads += allele1Reads;
	vc->alt_reads += allele2Reads;
      }

      // Associate zygosity,allele with the number of reads
      pair<pair<zygosity, string>, int> p;
      p.first = pair<zygosity, string>(zyg, vc->alleles[0]);
      p.second = allele1Reads;
      map<pair<zygosity, string>, int>::iterator readCountIter;
      readCountIter = vc->alleleReadCount.find(p.first);
      if (readCountIter == vc->alleleReadCount.end()) {
	vc->alleleReadCount.insert(p);
      } else {
	(*readCountIter).second += p.second;
      }

      p.first = pair<zygosity, string>(zyg, vc->alleles[1]);
      p.second = allele2Reads;
      readCountIter = vc->alleleReadCount.find(p.first);
      if (readCountIter == vc->alleleReadCount.end()) {
	vc->alleleReadCount.insert(p);
      } else {
	(*readCountIter).second += p.second;
      }

    }
  }
}
					 
void
variant_stats_for_location_range(GeneLocation range,
				 map<GeneLocation, TargetVariant*> variantMap,
				 vector<pair<GeneLocation, pair<int, int> > >& results_in_range) {

  map<GeneLocation, TargetVariant*>::iterator variantIterator;
  map<pair<zygosity, string>, int> alleleCounts;
  assert(results_in_range.size() == 0);

  for (variantIterator = variantMap.begin(); variantIterator != variantMap.end();
       ++variantIterator) {
    vector<VariantAndCounts*>& variantStats = (*variantIterator).second->variant_stats;
    if (variantStats.size() == 0) {
      continue;
    }

    int total_ref_reads = 0;
    int total_alt_reads = 0;

    GeneLocation g = (*variantIterator).first;
    if (!range.isContained(g)) {
      continue;
    }
    pair<GeneLocation, pair<int, int> > one_bucket_counts(g, pair<int, int>(0, 0));;

    vector<VariantAndCounts*>::iterator it;

    // Asserting this for now because we haven't tested it with more
    // than one VCF file, and supporting more than one VCF files turns
    // out to not be relevant to the analysis.
    //    assert (variantStats.size() == 1);

    for (it = variantStats.begin(); it != variantStats.end(); ++it) {
      VariantAndCounts* oneVcfStats = *it;
      total_alt_reads += oneVcfStats->alt_reads;
      total_ref_reads += oneVcfStats->ref_reads;

      map<genotype, int>::iterator genoIter;
      for (genoIter = oneVcfStats->genotypeCounts.begin();
	   genoIter != oneVcfStats->genotypeCounts.end();
	   ++genoIter) {
	pair<zygosity, string> p1, p2;
	zygosity z;
	genotype g = (*genoIter).first;
	// Extract indices into allele vector from genotype, which has
	// the format of M/N
	int allele1 = g[0] - 48;
	int allele2 = g[2] - 48;
	assert(g[1] == '/');
	if (allele1 == allele2) {
	  z = HOM;
	} else {
	  z = HET;
	}
	p1.first = z;
	p1.second = oneVcfStats->alleles[allele1];
	p2.first = z;
	p2.second = oneVcfStats->alleles[allele2];
	map<pair<zygosity, string>, int>::iterator alleleCountIter;

	alleleCountIter = alleleCounts.find(p1);
	// We need to add the number of reads for an allele with this genotype
	int readsForAllele = (*(oneVcfStats->alleleReadCount.find(p1))).second;
	if (alleleCountIter == alleleCounts.end()) {
	  alleleCounts.insert(*(new pair<pair<zygosity, string>, int>(p1, readsForAllele)));
	} else {
	  (*alleleCountIter).second += readsForAllele;
	}

	alleleCountIter = alleleCounts.find(p2);
	readsForAllele = (*(oneVcfStats->alleleReadCount.find(p2))).second;
	if (alleleCountIter == alleleCounts.end()) {
	  alleleCounts.insert(*(new pair<pair<zygosity, string>, int>(p2, readsForAllele)));
	} else {
	  (*alleleCountIter).second += readsForAllele;
	}
      }
    }
    // If both of these is 0, there were no heterozygous samples for
    // this variant.
    if (total_ref_reads != 0 && total_alt_reads != 0) {
      one_bucket_counts.second = pair<int, int>(total_ref_reads, total_alt_reads);
      results_in_range.push_back(one_bucket_counts);
    }
  }
  // map<pair<zygosity, string>, int>::iterator alleleCountIter;
  // for (alleleCountIter = alleleCounts.begin();
  //      alleleCountIter != alleleCounts.end();
  //      ++alleleCountIter) {
  //   pair<pair<zygosity, string>, int> p = *alleleCountIter;
  //   pair<zygosity, string> key = p.first;
  //   int count = p.second;
  //   cout << (key.first ? "HET" : "HOM")  << "/" << key.second << "\t" << count << endl;
  // }
}
				      
int main(int argc, char *argv[]) {
  map<GeneLocation, TargetVariant*> variantMap;
  if (argc < 2) {
    cout << "invalid usage" << endl;
    return 1;
  }

  // These are the indices into the argv[] vector for each parameter.
  int vcf_file_index = 1;
  int aggregation_file_index = 2;
  int chromosome_argument_index = 3;

  vector<GeneLocation> aggregation_buckets;
  read_aggregation_buckets(argv[aggregation_file_index], aggregation_buckets);

  string chromosome;
  if (argc  > 3) {
    chromosome = (argv[chromosome_argument_index]);
  } else {
    chromosome = "";
  }

  accumulate_read_counts_for_vcf_file(argv[vcf_file_index], variantMap, chromosome);

  cout << variantMap.size() << " AA-changing variants stored in memory." << endl;
  
  vector<GeneLocation>::iterator aggIter;

  for (aggIter = aggregation_buckets.begin(); 
       aggIter != aggregation_buckets.end();
       ++aggIter) {
    cout << "For aggregation bucket " << *aggIter << endl;

    vector<pair<GeneLocation, pair<int, int> > > bucket_aggregation;
    variant_stats_for_location_range(*aggIter, variantMap, bucket_aggregation);

    vector<pair<GeneLocation, pair<int, int> > >::iterator iter;
    for (iter = bucket_aggregation.begin(); iter != bucket_aggregation.end();
	 ++iter) {
      cout << "\t" << iter->first << ": " << iter->second.first << "/" << iter->second.second << endl;
    }
  }
}
